import requests
import json
import os
import time
import sys

# 获取API密钥并验证
gemini_api_key = os.getenv("GEMINI_API_KEY")
if not gemini_api_key:
    print("错误：GEMINI_API_KEY 环境变量未设置。")
    sys.exit(1)

# API配置
GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:streamGenerateContent?alt=sse"
HEADERS = {
    "x-goog-api-key": gemini_api_key,
    "Content-Type": "application/json"
}

def main():
    try:
        while True:
            user_input = input("请输入您的问题 (输入 '0' 退出): \n")
            if user_input.strip() == '0':
                print("程序已退出。")
                break

            # 构建请求体
            payload = {"contents": [{"parts": [{"text": user_input}]}]}
            
            try:
                with requests.post(
                    GEMINI_URL,
                    headers=HEADERS,
                    data=json.dumps(payload),
                    stream=True
                ) as resp:
                    resp.raise_for_status()
                    print("\nAI 助手回答:")
                    
                    # 流式处理响应
                    for line in resp.iter_lines():
                        if line and line.startswith(b'data:'):
                            try:
                                data = json.loads(line.decode('utf-8').lstrip('data:').strip())
                                for cand in data.get('candidates', []):
                                    for part in cand.get('content', {}).get('parts', []):
                                        text = part.get('text', '')
                                        if text:
                                            # 逐字打印
                                            for c in text:
                                                print(c, end='', flush=True)
                                                time.sleep(0.005)
                            except json.JSONDecodeError:
                                continue
                            except Exception as e:
                                print(f"\n处理错误: {e}", file=sys.stderr)
                    
                    print(f"\n{'-'*50}")

            except requests.exceptions.RequestException as e:
                print(f"\n请求错误: {e}", file=sys.stderr)

    finally:
        print("程序已安全退出。")

if __name__ == "__main__":
    main()